from __future__ import annotations

import datetime
import os
import sqlite3
from pathlib import Path
from typing import TYPE_CHECKING, Any, List

from gradio.flagging import FlaggingCallback
from gradio_client.documentation import document, set_documentation_group

if TYPE_CHECKING:
    from gradio.components import IOComponent

set_documentation_group("flagging")

LOG_TABLE_NAME = 'log'

CREATE_LOG_FORMAT_SQL = """
create table {TABLE_NAME}
(
    id        integer
        constraint log_pk
            primary key autoincrement,
    {COLUMNS}
    flag      integer not null,
    username  TEXT    not null,
    timestamp TEXT    not null
);
"""


@document()
class DBCSVLogger(FlaggingCallback):
    """
    The default implementation of the FlaggingCallback abstract class. Each flagged
    sample (both the input and output data) is logged to a CSV file with headers on the machine running the gradio app.
    Example:
        import gradio as gr
        def image_classifier(inp):
            return {'cat': 0.3, 'dog': 0.7}
        demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label",
                            flagging_callback=CSVLogger())
    Guides: using-flagging
    """

    def __init__(self):
        self.components = None
        self.flagging_dir = None

    def setup(
            self,
            components: List[IOComponent],
            flagging_dir: str | Path,
    ):
        self.components = components
        self.flagging_dir = flagging_dir
        os.makedirs(flagging_dir, exist_ok=True)

    def flag(
            self,
            flag_data: List[Any],
            flag_option: str = "",
            username: str | None = None,
    ) -> int:
        # HEADER / CREATE DATABASE & TABLE
        flagging_dir = self.flagging_dir
        log_filepath = Path(flagging_dir) / "log.sqlite"
        with sqlite3.connect(log_filepath) as conn:
            res = conn.execute(f"select sum(name = '{LOG_TABLE_NAME}') from sqlite_master where type = 'table';")
            component_names = [getattr(component, "label", None) or f"component {idx}"
                               for idx, component in enumerate(self.components)]
            if res.fetchone()[0] != 1:
                CREATE_LOG_SQL = CREATE_LOG_FORMAT_SQL.format_map({
                    'TABLE_NAME': LOG_TABLE_NAME,
                    'COLUMNS': "\n".join(
                        [f'`{component_name}` TEXT NOT NULL,' for component_name in component_names]
                    )
                })
                conn.execute(CREATE_LOG_SQL)

            csv_data = []
            for idx, (component, sample) in enumerate(zip(self.components, flag_data)):
                csv_data.append(str(sample))
            csv_data.append(flag_option)
            csv_data.append(username if username is not None else "")
            csv_data.append(str(datetime.datetime.now()))

            # INSERT DATA
            conn.execute(f"""
            insert into `{LOG_TABLE_NAME}` (`{'`,`'.join(component_names)}`, flag, username, timestamp) 
                values ({('?,'*(len(component_names)+3))[:-1]});""", csv_data)

            res = conn.execute("select count(id) from log;")
            line_count = res.fetchone()[0]
        return line_count
